# optimizer
# optimizer = dict(type='SGD', lr=0.005, momentum=0.9, weight_decay=0.0001)
# optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2))
# # learning policy
# lr_config = dict(
#     policy='step',
#     warmup='linear',
#     warmup_iters=500,
#     warmup_ratio=0.001,
#     step=[24, 33])
# total_epochs = 36


# optimizer for all other parts
optimizer = dict(
    type='SGD', 
    lr=0.005, 
    momentum=0.9, 
    weight_decay=0.0001,
    paramwise_cfg=dict(
        custom_keys={
            # Setting lr_mult for different parts
            'seg_head': dict(lr_mult=2.0, decay_mult=5.0),   # This should effectively double the learning rate for seg_head
        }
    )
)

optimizer_config = dict(
    grad_clip=dict(max_norm=35, norm_type=2)
)

# learning policy for segmentation head with "poly" policy
# seg_lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=True) 

# learning policy for other parts with "step" policy
# lr_config = dict(
#     policy='step',
#     warmup='linear',
#     warmup_iters=500,
#     warmup_ratio=0.001,
#     step=[24, 33]
# )

lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)

total_epochs = 36
